#!/usr/bin/env python
"""

    get_genomic_sequence.py
    [--log_file PATH]
    [--quiet]

"""

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   options        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

from optparse import OptionParser
import sys, os
import os.path
import StringIO

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    sys.path.append(os.path.abspath(os.path.join(exe_path,"..", "python_modules")))
    module_name = os.path.split(sys.argv[0])[1]
    module_name = os.path.splitext(module_name)[0];
else:
    module_name = __name__


parser = OptionParser(version="%prog 1.0", usage = "\n\n    %progs [options]")
parser.add_option("-c", "--contig", dest="contig",
                  metavar="CONTIG", 
                  type="string",
                  help="Name of contig to retrieve.")
parser.add_option("-p", "--sequence_position", dest="sequence_position",
                  metavar="INTEGER", 
                  type="int",
                  help="Position to start retrieving from.")
parser.add_option("-l", "--sequence_length", dest="sequence_length",
                  metavar="INTEGER", 
                  type="int",
                  help="length of sequence to retrieve.")
parser.add_option("-s", "--sequence_filename", dest="sequence_filename",
                    metavar="FILE", 
                    type="string",
                    help="Name and path of indexed FASTA file")
parser.add_option("-i", "--index_filename", dest="index_filename",
                    metavar="FILE", 
                    type="string",
                    help="Name and path of FASTA index.")
parser.add_option("-r", "--reverse_complement", dest="reverse_complement",
                    action="store_true", default=False,
                    help="Reverse complement sequence.")




#
#   general options: verbosity / logging
# 
parser.add_option("-v", "--verbose", dest = "verbose",
                  action="store_true", default=False,
                  help="Print progress to STDERR")
parser.add_option("-L", "--log_file", dest="log_file",
                  metavar="FILE", 
                  type="string",
                  help="Name and path of log file")
parser.add_option("--skip_parameter_logging", dest="skip_parameter_logging",
                    action="store_true", default=False,
                    help="Do not print program parameters to log.")
parser.add_option("--debug", dest="debug",
                    action="store_true", default=False,
                    help="Set default program parameters in debugging mode.")


if __name__ == '__main__':

    # get help string
    f =StringIO.StringIO()
    parser.print_help(f)
    helpstr = f.getvalue()
    (options, remaining_args) = parser.parse_args()
    
    if options.debug:
        options.log_file = os.path.join("get_genomic_sequence.log")
    
    # mandatory options
    from options import check_mandatory_options
    mandatory_options = [   "index_filename", "sequence_filename", 
                            "sequence_length", "sequence_position", "contig"]
    check_mandatory_options (options, mandatory_options, helpstr)
    
    
#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   imports        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import StringIO
import re
import operator

from collections import defaultdict, namedtuple
import logging
from lg_program_logging import setup_std_logging
from options import get_option_strings
from reverse_complement_dna import reverse_complement

contig_info = namedtuple('contig_info', 'seq_beg seq_len ignored acc_beg acc_len')


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
#_________________________________________________________________________________________
#
#       get_genomic_sequence_coordinates_from_panda
# 
#_________________________________________________________________________________________
def get_genomic_sequence_coordinates_from_panda (cursor ):
    """
    get cached indexed genomic data from panda postgresql database

    """
    import psycopg2
    import psycopg2.extensions
    cursor.execute("""
        SELECT species, name, seq_beg, seq_len, chrm_offset
                        FROM taxon.chromosomes;
        """)

    genomic_sequence_coodinates = defaultdict(dict)
    for (species, chromosome, seq_beg, seq_len, chrm_offset) in cursor.fetchall():
        # convert to numbers before saving
        seq_beg, seq_len, chrm_offset = map(int,  (seq_beg, seq_len, chrm_offset))
        genomic_sequence_coodinates[species][chromosome] = [seq_beg, seq_len, chrm_offset]

    return genomic_sequence_coodinates


#_________________________________________________________________________________________
#
#       get_genomic_sequence_coordinates_from_index_file
# 
#_________________________________________________________________________________________
def get_genomic_sequence_coordinates_from_index_file (index_file):
    """
    get cached indexed genomic data from index file

    """
    genomic_sequence_coordinates = dict()

    try:
        for i, line in enumerate(index_file):
            line = line.strip()
            #ignore empty or comment lines
            if not len(line)  or line[0] == '#':
                continue
            data = line.split("\t")
            if len(data) == 2:
                alias, contig = data
                if alias in genomic_sequence_coordinates:
                    genomic_sequence_coordinates[contig] = genomic_sequence_coordinates[alias]
                continue
            elif len(data) == 5:
                raise Exception("Can't read compressed files")
            elif len(data) != 4:
                raise Exception("Problems reading %s at line#%d" % (line, i))

            contig, acc_start, seq_start, seq_len  = data
                
            acc_start, seq_start, seq_len = map(int,  (acc_start, 
                                                                seq_start, seq_len))
            # ignore >
            acc_start += 1
            acc_beg = seq_start - acc_start -1
    
            
            # split ignoring ">" in chromosome name
            import re
            genomic_sequence_coordinates[contig] = contig_info(seq_start, seq_len, 0, acc_start, acc_beg)
            alt_contig = contig.split(":")[0]
            # chrom_offset no longer necessary and set to 0
            genomic_sequence_coordinates[alt_contig] = genomic_sequence_coordinates[contig]
    except:
        print >>sys.stderr, "Problems reading %s at line#%d" % (line, i)
        raise

    return genomic_sequence_coordinates


#
#   use cache
#
cached_fh = None
cached_filename = ""

#_________________________________________________________________________________________
#
#       get_genomic_sequence
# 
#_________________________________________________________________________________________
import mmap
import sys
def get_genomic_sequence (contig, beg, seq_len, genomic_sequence_coordinates, 
                            sequence_filename):
    """
        sequence_filename = "/some_dir/some_where/genomic.fa"
    """
    global cached_fh
    global cached_filename


    #
    #   open file and add to cache
    #
    if sequence_filename != cached_filename:
        if cached_fh:
           cached_fh.close()
           del cached_fh
        cached_fh = open(sequence_filename)
        cached_filename = sequence_filename
        try:
            file_mmap = mmap.mmap(cached_fh.fileno(), 0, mmap.MAP_PRIVATE, mmap.PROT_READ)
            cached_fh = file_mmap
        # not enough memory
        except EnvironmentError:
            pass

    #if contig[0:3] == "chr":
    #    contig = contig[3:]
    new_contig = None
    if not genomic_sequence_coordinates.has_key(contig):
        contig_sans_version = contig.split('.')[0]
        if genomic_sequence_coordinates.has_key(contig_sans_version):
            new_contig = contig = contig_sans_version
        else:
            #        5|NT_113890.1
            contig_sans_chromosome_prefix = contig_sans_version.split('|')
            if (len(contig_sans_chromosome_prefix) > 1 and
                genomic_sequence_coordinates.has_key(contig_sans_chromosome_prefix[1])):
                new_contig = contig = contig_sans_chromosome_prefix[1]
            else:
                raise RuntimeError, "File %s has no contig [%s] %s" % (cached_filename, contig, str(genomic_sequence_coordinates))

    chrm_beg, chrm_len, chrm_offset, acc_beg, acc_len = genomic_sequence_coordinates[contig]

    if beg < chrm_offset:
        raise RuntimeError, ("In contig [%s], %d is less than the first " 
                                        "available sequence position (%d)." %
                    (contig, beg, chrm_offset))

    if beg > chrm_len + chrm_offset:
        raise RuntimeError, ("In contig [%s], the starting position %d is larger than the contig length (%d)." %
                    (contig, beg, chrm_len + chrm_offset))

        

    if seq_len < 0:
        err_args = map(str, (contig, beg, seq_len,
                        cached_filename, chrm_beg, chrm_len, chrm_offset))
        sys.stderr.write("%s:%d" % (str( err_args), len(err_args)))
        format_str = "genomic sequence could not be read for negative sequence lengths.\n"+\
                    "Parameters = contig[%s] [%s+%s] from [%s]\n"+       \
                    "Contig params : [b = %s, l = %s, offset = %s]"
        raise RuntimeError, (format_str % tuple(err_args))

    if beg + seq_len > chrm_len + chrm_offset:
        seq_len = chrm_len + chrm_offset - beg
        #raise RuntimeError, ("In contig [%s], %s + %s is beyond "
        #            "than the available sequence length (%s)" %
        #            (contig, str(beg), str(seq_len), str(chrm_len + chrm_offset)))

    # 
    # read the file at the specied position
    #
    cached_fh.seek(chrm_beg + beg - chrm_offset)
    sequence = cached_fh.read(seq_len)

    if len(sequence) != seq_len:
        err_args = map(str, (len(sequence), contig, beg, seq_len,
                        cached_filename, chrm_beg, chrm_len, chrm_offset))
        sys.stderr.write("%s:%d" % (str( err_args), len(err_args)))
        format_str = "Only able to read (%s) of genomic sequence.\n"+\
                    "Parameters = contig[%s] [%s+%s] from [%s]\n"+       \
                    "Contig params : [b = %s, l = %s, offset = %s]"
        raise RuntimeError, (format_str % tuple(err_args))

    return sequence, new_contig


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Main logic


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888


if __name__ == '__main__':
    logger = logging.getLogger(module_name)
    setup_std_logging(logger, options.log_file, options.verbose)

    if not options.skip_parameter_logging:
        programme_name = os.path.split(sys.argv[0])[1]
        logger.info("%s %s" % (programme_name, " ".join(get_option_strings(parser, options))))

    genomic_sequence_coordinates = get_genomic_sequence_coordinates_from_index_file (open(options.index_filename))
    seq = get_genomic_sequence (options.contig, 
                                options.sequence_position, 
                                options.sequence_length, 
                                genomic_sequence_coordinates, 
                                options.sequence_filename)
    if options.reverse_complement:
        print reverse_complement(seq[0])
    else:
        print seq[0]

